# -*- coding:utf-8 -*-
'''
@paper:Noise Adaptive Fuzzy Switching Median Filter for Salt-and-Pepper Noise Reduction
'''
import numpy as np


# Adaptive Weighted Mean Filter
class Nafsmf(object):
    def __init__(self, h, s_max, T1, T2):
        self.h = h
        self.s_max = s_max
        self.T1 = T1
        self.T2 = T2

    def t_zero_pad_matrix(self, A, t):
        '''图像四周填充t层0'''
        m, n = np.shape(A)
        img_pad = np.zeros((m + t, n + t))
        img_pad[t:, t:] = A
        return img_pad

    def binary_noise_mask(self, A):
        '''噪声像素置0，其他像素置1'''
        B = np.zeros(np.shape(A))
        B[(A == 255) | (A == 0)] = 0
        B[(A != 255) & (A != 0)] = 1
        return B

    def k_approximate_matrix(self, img_pad, i, j, k, t):
        k_approximate_mat = img_pad[(i + t - k):(i + t + k + 1), (j + t - k):(j + t + k + 1)]
        return k_approximate_mat

    def fuzzy(self, D):
        if D < self.T1:
            F = 0
        elif D >= self.T2:
            F = 1
        else:
            F = (D - self.T1) / (self.T2 - self.T1)
        return F

    def process_image(self, X):
        m, n = np.shape(X)
        Y = X.copy()
        N = self.binary_noise_mask(X)
        img_pad = self.t_zero_pad_matrix(X, self.s_max)
        for i in range(m):
            for j in range(n):
                if N[i, j] == 0:
                    s = 1
                    while s <= self.s_max:
                        W = self.k_approximate_matrix(img_pad, i, j, k=s, t=self.s_max)
                        G = np.sum(self.binary_noise_mask(W))
                        if G >= 1:
                            M = np.median(W[self.binary_noise_mask(W).astype(bool)])
                        elif G == 0 and s == 3:
                            tem = self.k_approximate_matrix(img_pad, i, j, k=1, t=self.s_max)
                            tem = tem.flatten()
                            M = np.median(tem[0:4])
                        else:
                            s = s + 1
                            continue
                        tem = self.k_approximate_matrix(img_pad, i, j, k=1, t=self.s_max)
                        d = np.abs(tem - X[i, j])
                        D = np.max(d)
                        F = self.fuzzy(D)
                        Y[i, j] = (1 - F) * X[i, j] + F * M
                        break
        return Y
